'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 1_study_multy_query.py
* @Time: 2025/10/30
* @All Rights Reserve By Brtc
'''
import dotenv
import weaviate
from langchain.retrievers import MultiQueryRetriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_weaviate import WeaviateVectorStore

dotenv.load_dotenv()

#1、构建向量数据库
client = weaviate.connect_to_local("192.168.106.129", 8080)
db = WeaviateVectorStore(
    client,
    index_name="TestDemo",
    text_key="text",
    embedding=OpenAIEmbeddings(model="text-embedding-3-small")
)
retriever = db.as_retriever(search_type="mmr")

#2、创建多查询检索器
multi_query_retriever = MultiQueryRetriever.from_llm(
    retriever=retriever,
    llm=ChatOpenAI(model="gpt-4o-mini", temperature=0),
    prompt=ChatPromptTemplate.from_template(
        "你是一个AI语言模型助手。你的任务是生成给定用户问题的3个不同版本，以从向量数据库中检索相关文档。"
        "通过提供用户问题的多个视角，你的目标是帮助用户克服基于距离的相似性搜索的一些限制。"
        "请用换行符分隔这些替代问题。"
        "原始问题：{question}"
    )
)

#3、执行检索
docs = multi_query_retriever.invoke("关于LLMOps应用配置文档有那些？")
for doc in docs:
    print("==================================================")
    print(doc.page_content[:50])
client.close()
